package com.agent.server.sse;

import com.agent.server.bean.Question;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.MediaType;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;
import reactor.core.publisher.FluxSink;

import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
import java.util.function.Consumer;

public class LlmSseClient {
    private static final Logger logger = LoggerFactory.getLogger(LlmSseClient.class);
    private final ObjectMapper objectMapper = new ObjectMapper();
    private final WebClient webClient;
    private final String requestId;
    private final String baseUrl;
    private final Map<String, String> headers;

    /**
     * Creates an SSE client for communicating with LLM services
     *
     * @param baseUrl The base URL of the LLM service
     * @param headers HTTP headers to include in requests
     */
    public LlmSseClient(String baseUrl, Map<String, String> headers) {
        this.baseUrl = baseUrl;
        this.headers = headers;
        this.requestId = UUID.randomUUID().toString();
        this.webClient = WebClient.builder()
                .baseUrl(baseUrl)
                .build();
    }

    /**
     * Sends a request to the LLM service and returns a Flux of response events
     *
     * @param content The user prompt to send to the LLM
     * @return Flux of SSE response strings
     */
    public Flux<String> streamRequest(String content) {
        try {
            Map<String, Object> request = new HashMap<>();
            request.put("bot_id", "7501496633164120114");
            request.put("user_id", requestId);
            request.put("auto_save_history", false);
            request.put("stream", true);
            request.put("additional_messages", Collections.singletonList(new Question(content)));

            String jsonRequest = objectMapper.writeValueAsString(request);
            logger.info("Sending request to LLM: {}", jsonRequest);
            
            return webClient.post()
                    .uri(baseUrl)
                    .headers(httpHeaders -> {
                        headers.forEach(httpHeaders::add);
                    })
                    .contentType(MediaType.APPLICATION_JSON)
                    .body(BodyInserters.fromValue(jsonRequest))
                    .accept(MediaType.TEXT_EVENT_STREAM)
                    .retrieve()
                    .bodyToFlux(String.class)
                    .doOnNext(event -> logger.info("Received SSE event: {}", event))
                    .doOnError(error -> logger.error("Error in SSE stream", error))
                    .doOnComplete(() -> logger.info("SSE stream completed"));
        } catch (JsonProcessingException e) {
            logger.error("Error creating LLM request", e);
            return Flux.error(e);
        }
    }
    
    /**
     * Creates a Flux that can be consumed by a Consumer function
     * 
     * @param prompt The user prompt
     * @param messageConsumer Consumer that will handle each SSE message
     * @return A Flux of response events
     */
    public Flux<String> streamRequestWithConsumer(String prompt, Consumer<String> messageConsumer) {
        return Flux.create(emitter -> {
            streamRequest(prompt)
                .subscribe(
                    response -> {
                        try {
                            messageConsumer.accept(response);
                            emitter.next(response);
                        } catch (Exception e) {
                            logger.error("Error processing SSE response", e);
                        }
                    },
                    error -> {
                        logger.error("SSE stream error", error);
                        emitter.error(error);
                    },
                    () -> {
                        logger.info("SSE stream completed");
                        emitter.complete();
                    }
                );
        }, FluxSink.OverflowStrategy.BUFFER);
    }
} 